first pass with model

app.py

@@ -1,5 +1,102 @@
 import streamlit as st
+import numpy as np
+import PIL
+from PIL import Image
+from streamlit_image_select import image_select
+from ultralytics import YOLO
+import cv2
+import matplotlib.pyplot as plt
+import os
+import pathlib
+import PIL
+import PIL.Image
+import xml.etree.ElementTree as ET
+import pybboxes as pbx
+from pybboxes import BoundingBox
+from pathlib import Path
+import colorsys
+import random
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
-st.write(x, 'times 2 is', x * 2)
+####################################################
+# Support functions
+####################################################
+
+# helper function to generate random colors for class boxes
+def generate_label_colors(count):
+  colors = []
+  for c in range(count):
+    h,s,l = random.random(), 0.5 + random.random()/2.0, 0.4 + random.random()/5.0
+    r,g,b = [int(256*i) for i in colorsys.hls_to_rgb(h,l,s)]
+    colors.append(tuple([int(r), int(g), int(b)]))
+  return colors
+
+# helper function to run model inference
+def run_inference(model, img_paths):
+  return model.predict(img_paths)
+
+# helper function to process result and return image with bbox overlays
+def process_inference_result(result, class_colors):
+
+  # extract result objects
+  img = result.orig_img
+  dh, dw, _ = img.shape
+  boxes = result.boxes.xywhn.tolist()
+  labels = [int(label) for label in result.boxes.cls]
+  conf = [float(label) for label in result.boxes.conf]
+
+  # create image
+  for i, bbox in enumerate(boxes):
+    x = bbox[0]
+    y = bbox[1]
+    w = bbox[2]
+    h = bbox[3]
+    voc_box = pbx.convert_bbox([x, y, w, h], from_type="yolo", to_type="voc", image_size=(dw, dh))
+    voc_x1 = voc_box[0]
+    voc_y1 = voc_box[1]
+    voc_x2 = voc_box[2]
+    voc_y2 = voc_box[3]
+    cv2.rectangle(img, (voc_x1, voc_y1), (voc_x2, voc_y2), class_colors[labels[i]], 2)
+    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)
+
+  return img
+
+def run_process_show(img_path):
+  results = model(img_path)
+  processed_image = process_inference_result(results[0], rand_class_colors)
+  return processed_image
+
+####################################################
+# Setup model and class parameters
+####################################################
+
+# init model
+model = YOLO("weights/best.pt")
+
+# setup label classes
+classes = ['A6', 'A17', 'A16', 'A15', 'A5', 'A20', 'A14', 'A12', 'A8', 'A2', 'A7', 'A18', 'A13', 'A4', 'A19', 'A1', 'A3', 'A10', 'A11', 'A9']
+
+# setup mapping of class labels to real aircraft names
+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']
+aircraft_lookup = {}
+for i in range(len(classes)):
+    aircraft_lookup['A' + str(i+1)] = aircraft_names[i]
+
+# generate bbox colors for each class
+rand_class_colors = generate_label_colors(len(classes))
+
+
+####################################################
+# Main UX Loop
+####################################################
+img = image_select(
+    label="Select an airbase",
+    images=[
+        cv2.imread("images/edwards.jpg"),
+        cv2.imread("images/edwards2.jpg"),
+    ],
+    captions=["Edwards AFB1", "Edwards AFB2"],
+)
+
+# process image through detector
+img2 = run_process_show(img)
+st.image(img2)

weights/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7aba73b51184250fe327b738f1ccaa3929ae75415356d5aaaafa4b191e0a05d
+size 6258073