SaladSlayer00 commited on
Commit
d7feb62
1 Parent(s): 18d3941
Files changed (2) hide show
  1. app.py +34 -0
  2. requirements.txt +5 -0
app.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import DetrImageProcessor, DetrForObjectDetection
3
+ from PIL import Image
4
+ import torch
5
+ import cv2
6
+ import numpy as np
7
+
8
+ def process_image(input_image):
9
+ processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
10
+ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
11
+ yellow = (0, 255, 255) # BGR
12
+ font = cv2.FONT_HERSHEY_SIMPLEX
13
+ stroke = 2
14
+
15
+ # Convert PIL image to OpenCV format
16
+ img = np.array(input_image)
17
+ img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
18
+
19
+ # Process the image
20
+ inputs = processor(images=input_image, return_tensors="pt")
21
+ outputs = model(**inputs)
22
+ target_sizes = torch.tensor([input_image.size[::-1]])
23
+ results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
24
+
25
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
26
+ cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), yellow, stroke)
27
+ cv2.putText(img, model.config.id2label[label.item()], (int(box[0]), int(box[1]-10)), font, 1, yellow, stroke, cv2.LINE_AA)
28
+
29
+ # Convert back to PIL image
30
+ return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
31
+
32
+ # Create Gradio interface
33
+ iface = gr.Interface(fn=process_image, inputs=gr.inputs.Image(), outputs="image")
34
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ timm
4
+ opencv-python
5
+ gradio