ThirdEyeData/Object-Detection-Using-FRCNN

1import streamlit as st2import torch3import torchvision4import torchvision.transforms as transforms5from torchvision.models.detection.faster_rcnn import FastRCNNPredictor6from torchvision.transforms import ToTensor7from PIL import Image, ImageDraw8import cv29import numpy as np10import pandas as pd11import os12 13import tempfile14from tempfile import NamedTemporaryFile15 16# Create an FRCNN model instance with the same structure as the saved model17model = torchvision.models.detection.fasterrcnn_resnet50_fpn(num_classes=91)18 19device = torch.device("cuda" if torch.cuda.is_available() else "cpu")20 21# Load the saved parameters into the model22model.load_state_dict(torch.load("frcnn_model.pth"))23 24# Define the classes for object detection25classes = [26            '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',27            'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign',28            'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',29            'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A',30            'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',31            'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',32            'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork',33            'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',34            'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',35            'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A',36            'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',37            'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase',38            'scissors', 'teddy bear', 'hair drier', 'toothbrush'39          ]40 41# Set the threshold for object detection. It is IoU (Intersection over Union)42threshold = 0.543 44st.title(""" Object Detection Using Faster-RCNN """)45 46# st.subheader("Prediction of Object Detection")47 48st.write(""" The Faster R-CNN (Region-based Convolutional Neural Network) is a powerful object detection model that combines deep 49             learning with region proposal networks to achieve highly accurate object detection in images. 50             It is trained on a large dataset of images and can detect a wide range of objects with high precision and recall. 51             The model is based on the ResNet-50 architecture, which allows it to capture complex visual features from the input image. 52             It uses a two-stage approach, first proposing regions of interest (RoIs) in the image and then classifying and refining the 53             object boundaries within these RoIs. This approach makes it extremely efficient and accurate in detecting multiple objects 54             in a single image.55         """)56 57images = ["test2.jpg","img7.jpg","img20.jpg","img23.jpg","test1.jpg","img18.jpg","img3.jpg","img15.jpg","img17.jpg"]58with st.sidebar:59    st.write("Choose an Image")60    st.image(images)61 62# define the function to perform object detection on an image63def detect_objects(image_path):64    # load the image65    image = Image.open(image_path).convert('RGB')66    67    # convert the image to a tensor68    image_tensor = ToTensor()(image).to(device)69    70    # run the image through the model to get the predictions71    model.eval()72    with torch.no_grad():73        predictions = model([image_tensor])74 75    # filter out the predictions below the threshold76    scores = predictions[0]['scores'].cpu().numpy()77    boxes = predictions[0]['boxes'].cpu().numpy()78    labels = predictions[0]['labels'].cpu().numpy()79    mask = scores > threshold80    scores = scores[mask]81    boxes = boxes[mask]82    labels = labels[mask]83 84    # create a new image with the predicted objects outlined in rectangles85    draw = ImageDraw.Draw(image)86    for box, label in zip(boxes, labels):87        88        # draw the rectangle around the object89        draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline='red')90        91        # write the object class above the rectangle92        class_name = classes[label]93        draw.text((box[0], box[1]), class_name, fill='yellow')94    95    # show the image96    st.write("Obects detected in the image are: ")97    st.image(image, use_column_width=True)98    # st.image.show()99 100 101file = st.file_uploader('Upload an Image', type=(["jpeg", "jpg", "png"]))102 103 104if file is None:105    st.write("Please upload an image file")106else:107    image = Image.open(file)108    st.write("Input Image")109    st.image(image, use_column_width=True)110    with NamedTemporaryFile(dir='.', suffix='.') as f:111        f.write(file.getbuffer())112        # your_function_which_takes_a_path(f.name)113        detect_objects(f.name)114 115# if file is None:116#     st.write("Please upload an image file")117# else:118#     image = Image.open(file)119#     st.write("Input Image")120#     st.image(image, use_column_width=True)121#     with NamedTemporaryFile(dir='.', suffix='.jpeg') as f:  # this line gives error and only accepts .jpeg and so used above snippet122#         f.write(file.getbuffer())                           # which will accepts all formats of images.123#         # your_function_which_takes_a_path(f.name)124#         detect_objects(f.name)125 126st.write(""" This Streamlit app provides a user-friendly interface for uploading an image and visualizing the output of the Faster R-CNN 127             model. It displays the uploaded image along with the predicted objects highlighted with bounding box overlays. The app allows128             users to explore the detected objects in the image, providing valuable insights and understanding of the model's predictions. 129             It can be used for a wide range of applications, such as object recognition, image analysis, and visual storytelling. 130             Whether it's identifying objects in real-world images or understanding the capabilities of state-of-the-art object detection131             models, this Streamlit app powered by Faster R-CNN is a powerful tool for computer vision tasks.132        """)133 134