first pass with model
Browse files- app.py +100 -3
- images/.DS_Store +0 -0
- images/edwards.jpg +0 -0
- images/edwards2.jpg +0 -0
- weights/best.pt +3 -0
app.py
CHANGED
@@ -1,5 +1,102 @@
|
|
1 |
import streamlit as st
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
|
3 |
-
|
4 |
-
|
5 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import streamlit as st
|
2 |
+
import numpy as np
|
3 |
+
import PIL
|
4 |
+
from PIL import Image
|
5 |
+
from streamlit_image_select import image_select
|
6 |
+
from ultralytics import YOLO
|
7 |
+
import cv2
|
8 |
+
import matplotlib.pyplot as plt
|
9 |
+
import os
|
10 |
+
import pathlib
|
11 |
+
import PIL
|
12 |
+
import PIL.Image
|
13 |
+
import xml.etree.ElementTree as ET
|
14 |
+
import pybboxes as pbx
|
15 |
+
from pybboxes import BoundingBox
|
16 |
+
from pathlib import Path
|
17 |
+
import colorsys
|
18 |
+
import random
|
19 |
|
20 |
+
####################################################
|
21 |
+
# Support functions
|
22 |
+
####################################################
|
23 |
+
|
24 |
+
# helper function to generate random colors for class boxes
|
25 |
+
def generate_label_colors(count):
|
26 |
+
colors = []
|
27 |
+
for c in range(count):
|
28 |
+
h,s,l = random.random(), 0.5 + random.random()/2.0, 0.4 + random.random()/5.0
|
29 |
+
r,g,b = [int(256*i) for i in colorsys.hls_to_rgb(h,l,s)]
|
30 |
+
colors.append(tuple([int(r), int(g), int(b)]))
|
31 |
+
return colors
|
32 |
+
|
33 |
+
# helper function to run model inference
|
34 |
+
def run_inference(model, img_paths):
|
35 |
+
return model.predict(img_paths)
|
36 |
+
|
37 |
+
# helper function to process result and return image with bbox overlays
|
38 |
+
def process_inference_result(result, class_colors):
|
39 |
+
|
40 |
+
# extract result objects
|
41 |
+
img = result.orig_img
|
42 |
+
dh, dw, _ = img.shape
|
43 |
+
boxes = result.boxes.xywhn.tolist()
|
44 |
+
labels = [int(label) for label in result.boxes.cls]
|
45 |
+
conf = [float(label) for label in result.boxes.conf]
|
46 |
+
|
47 |
+
# create image
|
48 |
+
for i, bbox in enumerate(boxes):
|
49 |
+
x = bbox[0]
|
50 |
+
y = bbox[1]
|
51 |
+
w = bbox[2]
|
52 |
+
h = bbox[3]
|
53 |
+
voc_box = pbx.convert_bbox([x, y, w, h], from_type="yolo", to_type="voc", image_size=(dw, dh))
|
54 |
+
voc_x1 = voc_box[0]
|
55 |
+
voc_y1 = voc_box[1]
|
56 |
+
voc_x2 = voc_box[2]
|
57 |
+
voc_y2 = voc_box[3]
|
58 |
+
cv2.rectangle(img, (voc_x1, voc_y1), (voc_x2, voc_y2), class_colors[labels[i]], 2)
|
59 |
+
cv2.putText(img, aircraft_lookup[classes[labels[i]]] + ' ' + str(round(conf[i], 2)), (voc_x1, voc_y1-5), cv2.FONT_HERSHEY_SIMPLEX, 1, class_colors[labels[i]], 2)
|
60 |
+
|
61 |
+
return img
|
62 |
+
|
63 |
+
def run_process_show(img_path):
|
64 |
+
results = model(img_path)
|
65 |
+
processed_image = process_inference_result(results[0], rand_class_colors)
|
66 |
+
return processed_image
|
67 |
+
|
68 |
+
####################################################
|
69 |
+
# Setup model and class parameters
|
70 |
+
####################################################
|
71 |
+
|
72 |
+
# init model
|
73 |
+
model = YOLO("weights/best.pt")
|
74 |
+
|
75 |
+
# setup label classes
|
76 |
+
classes = ['A6', 'A17', 'A16', 'A15', 'A5', 'A20', 'A14', 'A12', 'A8', 'A2', 'A7', 'A18', 'A13', 'A4', 'A19', 'A1', 'A3', 'A10', 'A11', 'A9']
|
77 |
+
|
78 |
+
# setup mapping of class labels to real aircraft names
|
79 |
+
aircraft_names = ['SU-35', 'C-130', 'C-17', 'C-5', 'F-16', 'TU-160', 'E-3', 'B-52', 'P-3C', 'B-1B', 'E-8', 'TU-22', 'F-15', 'KC-135', 'F-22', 'FA-18', 'TU-95', 'KC-10', 'SU-34', 'SU-24']
|
80 |
+
aircraft_lookup = {}
|
81 |
+
for i in range(len(classes)):
|
82 |
+
aircraft_lookup['A' + str(i+1)] = aircraft_names[i]
|
83 |
+
|
84 |
+
# generate bbox colors for each class
|
85 |
+
rand_class_colors = generate_label_colors(len(classes))
|
86 |
+
|
87 |
+
|
88 |
+
####################################################
|
89 |
+
# Main UX Loop
|
90 |
+
####################################################
|
91 |
+
img = image_select(
|
92 |
+
label="Select an airbase",
|
93 |
+
images=[
|
94 |
+
cv2.imread("images/edwards.jpg"),
|
95 |
+
cv2.imread("images/edwards2.jpg"),
|
96 |
+
],
|
97 |
+
captions=["Edwards AFB1", "Edwards AFB2"],
|
98 |
+
)
|
99 |
+
|
100 |
+
# process image through detector
|
101 |
+
img2 = run_process_show(img)
|
102 |
+
st.image(img2)
|
images/.DS_Store
ADDED
Binary file (6.15 kB). View file
|
|
images/edwards.jpg
ADDED
images/edwards2.jpg
ADDED
weights/best.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:c7aba73b51184250fe327b738f1ccaa3929ae75415356d5aaaafa4b191e0a05d
|
3 |
+
size 6258073
|