Spaces:
Runtime error
Runtime error
| 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. | |
| """) |