Spaces:
Runtime error
Runtime error
Commit
·
0463ddb
1
Parent(s):
6c85af3
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,133 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import torch
|
3 |
+
import torchvision
|
4 |
+
import torchvision.transforms as transforms
|
5 |
+
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
|
6 |
+
from torchvision.transforms import ToTensor
|
7 |
+
from PIL import Image, ImageDraw
|
8 |
+
import cv2
|
9 |
+
import numpy as np
|
10 |
+
import pandas as pd
|
11 |
+
import os
|
12 |
+
|
13 |
+
import tempfile
|
14 |
+
from tempfile import NamedTemporaryFile
|
15 |
+
|
16 |
+
# Create an FRCNN model instance with the same structure as the saved model
|
17 |
+
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(num_classes=91)
|
18 |
+
|
19 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
20 |
+
|
21 |
+
# Load the saved parameters into the model
|
22 |
+
model.load_state_dict(torch.load("frcnn_model.pth"))
|
23 |
+
|
24 |
+
# Define the classes for object detection
|
25 |
+
classes = [
|
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)
|
42 |
+
threshold = 0.5
|
43 |
+
|
44 |
+
st.title(""" Image Object Detections """)
|
45 |
+
|
46 |
+
# st.subheader("Prediction of Object Detection")
|
47 |
+
|
48 |
+
st.write(""" The Faster R-CNN (Region-based Convolutional Neural Network) is a cutting-edge 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 |
+
|
57 |
+
images = ["test2.jpg","img7.jpg","img20.jpg","img23.jpg","test1.jpg","img18.jpg"]
|
58 |
+
with st.sidebar:
|
59 |
+
st.write("Choose an Image")
|
60 |
+
st.image(images)
|
61 |
+
|
62 |
+
# define the function to perform object detection on an image
|
63 |
+
def detect_objects(image_path):
|
64 |
+
# load the image
|
65 |
+
image = Image.open(image_path).convert('RGB')
|
66 |
+
|
67 |
+
# convert the image to a tensor
|
68 |
+
image_tensor = ToTensor()(image).to(device)
|
69 |
+
|
70 |
+
# run the image through the model to get the predictions
|
71 |
+
model.eval()
|
72 |
+
with torch.no_grad():
|
73 |
+
predictions = model([image_tensor])
|
74 |
+
|
75 |
+
# filter out the predictions below the threshold
|
76 |
+
scores = predictions[0]['scores'].cpu().numpy()
|
77 |
+
boxes = predictions[0]['boxes'].cpu().numpy()
|
78 |
+
labels = predictions[0]['labels'].cpu().numpy()
|
79 |
+
mask = scores > threshold
|
80 |
+
scores = scores[mask]
|
81 |
+
boxes = boxes[mask]
|
82 |
+
labels = labels[mask]
|
83 |
+
|
84 |
+
# create a new image with the predicted objects outlined in rectangles
|
85 |
+
draw = ImageDraw.Draw(image)
|
86 |
+
for box, label in zip(boxes, labels):
|
87 |
+
|
88 |
+
# draw the rectangle around the object
|
89 |
+
draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline='red')
|
90 |
+
|
91 |
+
# write the object class above the rectangle
|
92 |
+
class_name = classes[label]
|
93 |
+
draw.text((box[0], box[1]), class_name, fill='yellow')
|
94 |
+
|
95 |
+
# show the image
|
96 |
+
st.write("Obects detected in the image are: ")
|
97 |
+
st.image(image, use_column_width=True)
|
98 |
+
# st.image.show()
|
99 |
+
|
100 |
+
|
101 |
+
file = st.file_uploader('Upload an Image', type=(["jpeg", "jpg", "png"]))
|
102 |
+
|
103 |
+
|
104 |
+
if file is None:
|
105 |
+
st.write("Please upload an image file")
|
106 |
+
else:
|
107 |
+
image = Image.open(file)
|
108 |
+
st.write("Input Image")
|
109 |
+
st.image(image, use_column_width=True)
|
110 |
+
with NamedTemporaryFile(dir='.', suffix='.' + file.name.split('.')[-1]) 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 snippet
|
122 |
+
# 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 |
+
|
126 |
+
st.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 allows
|
128 |
+
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 detection
|
131 |
+
models, this Streamlit app powered by Faster R-CNN is a powerful tool for computer vision tasks.
|
132 |
+
""")
|
133 |
+
|