import streamlit as st
import torch
import torchvision
import torchvision.transforms as transforms
from torchvision import datasets, models
from torchvision.transforms import functional as FT
from torchvision import transforms as T
from torch import nn, optim
from torch.nn import functional as F
from torch.utils.data import DataLoader, sampler, random_split, Dataset
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.transforms import ToTensor
from PIL import Image, ImageDraw
from pycocotools.coco import COCO
import cv2
import numpy as np
import pandas as pd
import os


import tempfile
from tempfile import NamedTemporaryFile

dataset_path = "Dataset"

#load classes
coco = COCO(os.path.join(dataset_path, "train", "_annotations.coco.json"))
categories = coco.cats
n_classes = len(categories.keys())

# load the faster rcnn model
modeltest = models.detection.fasterrcnn_mobilenet_v3_large_fpn(num_classes=4)
in_features = modeltest.roi_heads.box_predictor.cls_score.in_features # we need to change the head
modeltest.roi_heads.box_predictor = models.detection.faster_rcnn.FastRCNNPredictor(in_features, n_classes)

# Load the saved parameters into the model
modeltest.load_state_dict(torch.load("FRCNN_MODEL_3Classes_100Epochs.pth"))

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
modeltest.to(device)

# Number of classes
classes = ['pole', 'cross_arm', 'pole', 'tag']

st.title(""" Object Detection Using Faster-RCNN For Electrical Domain """)

# st.subheader("Prediction of Object Detection")

images = ["img16.jpg","img1.jpg","img2.jpg","img3.jpg","img4.jpg","img5.jpg","img6.jpg","img8.jpg",
          "img10.jpg","img11.jpg","img12.jpg","img13.jpg","img14.jpg","img15.jpg","img9.jpg"]

with st.sidebar:
    st.write("Choose an Image from Sample Images ")
    st.image(images)

# with st.sidebar:
#     st.write("Choose an Image From The DropDown")
#     selected_image = st.selectbox("Select an image", images)

    
# with st.sidebar:
#     st.write("Choose an Image")
#     for image in images:
#         with Image.open(image) as img:
#             st.image(img, width=100, quality=90)   # quality parameter is not there in image, it will give error
    
# with st.sidebar:
#     st.write("Choose an Image")
#     st.image(images,width=100)
    

# define the function to perform object detection on an image
def detect_objects(image_path):
    # load the image
    image = Image.open(image_path).convert('RGB')
    
    # convert the image to a tensor
    image_tensor = ToTensor()(image).to(device)
    
    # run the image through the model to get the predictions
    modeltest.eval()
    with torch.no_grad():
        predictions = modeltest([image_tensor])

    # filter out the predictions below the threshold
    threshold = 0.5
    scores = predictions[0]['scores'].cpu().numpy()
    boxes = predictions[0]['boxes'].cpu().numpy()
    labels = predictions[0]['labels'].cpu().numpy()
    mask = scores > threshold
    scores = scores[mask]
    boxes = boxes[mask]
    labels = labels[mask]

    # create a new image with the predicted objects outlined in rectangles
    draw = ImageDraw.Draw(image)
    for box, label in zip(boxes, labels):
        
        # draw the rectangle around the object
        draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline='red')
        
        # write the object class above the rectangle
        class_name = classes[label]
        draw.text((box[0], box[1]), class_name, fill='yellow')
    
    # show the image
    st.write("Obects detected in the image are: ")
    st.image(image, use_column_width=True)
    # st.image.show()

file = st.file_uploader('Upload an Image', type=(["jpeg", "jpg", "png"]))


if file is None:
    st.write("Please upload an image file")
else:
    image = Image.open(file)
    st.write("Input Image")
    st.image(image, use_column_width=True)
    with NamedTemporaryFile(dir='.', suffix='.') as f:
        f.write(file.getbuffer())
        # your_function_which_takes_a_path(f.name)
        detect_objects(f.name)

st.subheader("Model Description : ")
st.write(""" The Faster R-CNN model with MobileNet V3 Large as the backbone and Feature Pyramid Network (FPN) architecture is a popular 
             object detection model that combines high detection accuracy with efficient computation. The MobileNet V3 Large backbone 
             is a lightweight neural network architecture that reduces the number of parameters while maintaining high accuracy, 
             making it suitable for mobile and embedded devices. The FPN architecture enhances the feature representation of the model 
             by aggregating features from multiple scales and improving spatial resolution. This combination of a lightweight backbone 
             with an efficient feature extraction architecture makes Faster R-CNN with MobileNet V3 Large FPN a popular choice for 
             object detection in real-time applications and on devices with limited computational resources.
         """)