File size: 5,039 Bytes
4a75cb7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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.
""") |