PUSHPENDAR/SHIP_FASTERRCNN
0
1# import gradio as gr2# import cv23# import numpy as np4# from detectron2.config import get_cfg5# from detectron2.engine import DefaultPredictor6# from detectron2.utils.visualizer import Visualizer, ColorMode7# from detectron2.data import MetadataCatalog8# from huggingface_hub import hf_hub_download9# import os10 11# REPO_ID = os.getenv("MODEL_REPO_ID", "PUSHPENDAR/hrsid-ship-detection")12 13# os.makedirs("/app/hf_cache", exist_ok=True)14 15# print("Downloading model files...")16# MODEL_PATH = hf_hub_download(repo_id=REPO_ID, filename="model_final.pth", cache_dir="/app/hf_cache")17# CONFIG_PATH = hf_hub_download(repo_id=REPO_ID, filename="config.yaml", cache_dir="/app/hf_cache")18# print(f"Model: {MODEL_PATH} ✅")19# print(f"Config: {CONFIG_PATH} ✅")20 21# print("Loading Faster R-CNN model...")22# cfg = get_cfg()23# cfg.merge_from_file(CONFIG_PATH)24# cfg.MODEL.WEIGHTS = MODEL_PATH25# cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.526# cfg.MODEL.DEVICE = "cpu"27 28# MetadataCatalog.get("__unused").set(thing_classes=["ship"])29# predictor = DefaultPredictor(cfg)30# print("Model loaded ✅")31 32 33# def detect_ships(image, confidence_threshold):34# if image is None:35# return None, "Please upload an image."36 37# cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold38# img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)39# outputs = predictor(img_bgr)40# instances = outputs["instances"].to("cpu")41# keep = instances.scores >= confidence_threshold42# instances = instances[keep]43 44# metadata = MetadataCatalog.get("__unused")45# v = Visualizer(img_bgr[:, :, ::-1], metadata=metadata, scale=1.0, instance_mode=ColorMode.IMAGE)46# out = v.draw_instance_predictions(instances)47 48# result_img = out.get_image()49# num_ships = len(instances)50# scores = instances.scores.tolist()51 52# info = f"✅ Detected {num_ships} ship(s)\n"53# if scores:54# info += "Confidence scores: " + ", ".join([f"{s:.2f}" for s in scores])55# if hasattr(instances, "pred_boxes"):56# boxes = instances.pred_boxes.tensor.tolist()57# info += "\n\nBounding boxes (x1,y1,x2,y2):\n"58# for i, (box, score) in enumerate(zip(boxes, scores)):59# x1, y1, x2, y2 = [int(v) for v in box]60# info += f" Ship {i+1}: [{x1},{y1},{x2},{y2}] conf={score:.2f}\n"61# else:62# info += "No ships detected above threshold."63 64# return result_img, info65 66 67# with gr.Blocks(title="🚢 HRSID Ship Detection") as demo:68# gr.Markdown("# 🚢 HRSID Ship Detection")69# gr.Markdown("Upload a SAR image to detect ships using Faster R-CNN with ResNet-101, trained on HRSID dataset.")70# with gr.Row():71# with gr.Column():72# image_input = gr.Image(type="pil", label="Upload SAR Image")73# threshold = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Confidence Threshold")74# btn = gr.Button("Detect Ships", variant="primary")75# with gr.Column():76# image_output = gr.Image(type="numpy", label="Detection Result")77# info_output = gr.Textbox(label="Detection Info", lines=10)78 79# btn.click(fn=detect_ships, inputs=[image_input, threshold], outputs=[image_output, info_output])80 81# if __name__ == "__main__":82# demo.launch(server_name="0.0.0.0", server_port=7860)83import os84import tempfile85from copy import deepcopy86 87import cv288import gradio as gr89import numpy as np90from detectron2.config import get_cfg91from detectron2.data import MetadataCatalog92from detectron2.engine import DefaultPredictor93from detectron2.utils.visualizer import ColorMode, Visualizer94from huggingface_hub import hf_hub_download95 96# ── Model loading ────────────────────────────────────────────────────────────97 98REPO_ID = os.getenv("MODEL_REPO_ID", "PUSHPENDAR/hrsid-ship-detection")99 100os.makedirs("/app/hf_cache", exist_ok=True)101 102print("Downloading model files...")103MODEL_PATH = hf_hub_download(104 repo_id=REPO_ID,105 filename="model_final.pth",106 cache_dir="/app/hf_cache",107 token=os.getenv("HF_TOKEN"), # uses secret if set, else None (public repos)108)109CONFIG_PATH = hf_hub_download(110 repo_id=REPO_ID,111 filename="config.yaml",112 cache_dir="/app/hf_cache",113 token=os.getenv("HF_TOKEN"),114)115print(f"Model: {MODEL_PATH} ✅")116print(f"Config: {CONFIG_PATH} ✅")117 118print("Loading Faster R-CNN model...")119_base_cfg = get_cfg()120_base_cfg.merge_from_file(CONFIG_PATH)121_base_cfg.MODEL.WEIGHTS = MODEL_PATH122_base_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5123_base_cfg.MODEL.DEVICE = "cpu"124_base_cfg.freeze() # make it immutable so we always deepcopy before mutating125 126MetadataCatalog.get("__unused").set(thing_classes=["ship"])127print("Model loaded ✅")128 129 130# ── Helpers ──────────────────────────────────────────────────────────────────131 132def get_predictor(confidence_threshold: float) -> DefaultPredictor:133 """Return a fresh predictor with the requested threshold.134 deepcopy avoids mutating the global frozen cfg across concurrent requests.135 """136 cfg = deepcopy(_base_cfg)137 cfg.defrost()138 cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = confidence_threshold139 return DefaultPredictor(cfg)140 141 142def run_inference(img_bgr: np.ndarray, confidence_threshold: float):143 """Run detection on a single BGR frame. Returns (result_bgr, instances)."""144 predictor = get_predictor(confidence_threshold)145 outputs = predictor(img_bgr)146 instances = outputs["instances"].to("cpu")147 instances = instances[instances.scores >= confidence_threshold]148 149 metadata = MetadataCatalog.get("__unused")150 v = Visualizer(151 img_bgr[:, :, ::-1],152 metadata=metadata,153 scale=1.0,154 instance_mode=ColorMode.IMAGE,155 )156 out = v.draw_instance_predictions(instances)157 result_rgb = out.get_image() # H×W×3 RGB158 result_bgr = cv2.cvtColor(result_rgb, cv2.COLOR_RGB2BGR)159 return result_bgr, instances160 161 162def build_info(instances) -> str:163 num = len(instances)164 scores = instances.scores.tolist()165 info = f"✅ Detected {num} ship(s)\n"166 if scores:167 info += "Confidence scores: " + ", ".join([f"{s:.2f}" for s in scores])168 if hasattr(instances, "pred_boxes"):169 boxes = instances.pred_boxes.tensor.tolist()170 info += "\n\nBounding boxes (x1,y1,x2,y2):\n"171 for i, (box, score) in enumerate(zip(boxes, scores)):172 x1, y1, x2, y2 = [int(c) for c in box]173 info += f" Ship {i+1}: [{x1},{y1},{x2},{y2}] conf={score:.2f}\n"174 else:175 info += "No ships detected above threshold."176 return info177 178 179# ── Image tab ────────────────────────────────────────────────────────────────180 181def detect_ships_image(image, confidence_threshold):182 if image is None:183 return None, "Please upload an image."184 img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)185 result_bgr, inst = run_inference(img_bgr, confidence_threshold)186 result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)187 return result_rgb, build_info(inst)188 189 190# ── Video tab ────────────────────────────────────────────────────────────────191 192def detect_ships_video(video_path, confidence_threshold, progress=gr.Progress()):193 if video_path is None:194 return None, "Please upload a video."195 196 cap = cv2.VideoCapture(video_path)197 if not cap.isOpened():198 return None, "Could not open video file."199 200 total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))201 fps = cap.get(cv2.CAP_PROP_FPS) or 25202 w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))203 h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))204 205 out_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)206 out_path = out_file.name207 out_file.close()208 209 fourcc = cv2.VideoWriter_fourcc(*"mp4v")210 writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h))211 212 frame_idx = 0213 total_ships = 0214 max_per_frame = 0215 216 while True:217 ret, frame = cap.read()218 if not ret:219 break220 221 result_bgr, inst = run_inference(frame, confidence_threshold)222 writer.write(result_bgr)223 224 n = len(inst)225 total_ships += n226 max_per_frame = max(max_per_frame, n)227 frame_idx += 1228 229 if total_frames > 0:230 progress(231 frame_idx / total_frames,232 desc=f"Processing frame {frame_idx}/{total_frames}",233 )234 235 cap.release()236 writer.release()237 238 info = (239 f"✅ Video processed: {frame_idx} frames\n"240 f"Total ship detections across all frames: {total_ships}\n"241 f"Peak ships in a single frame: {max_per_frame}\n"242 f"FPS: {fps:.1f} | Resolution: {w}×{h}"243 )244 return out_path, info245 246 247# ── UI ───────────────────────────────────────────────────────────────────────248 249with gr.Blocks(title="🚢 HRSID Ship Detection") as demo:250 gr.Markdown("# 🚢 HRSID Ship Detection")251 gr.Markdown(252 "Detect ships in SAR images **or videos** using "253 "Faster R-CNN with ResNet-101, trained on the HRSID dataset."254 )255 256 with gr.Tabs():257 258 with gr.Tab("🖼️ Image Detection"):259 with gr.Row():260 with gr.Column():261 img_input = gr.Image(type="pil", label="Upload SAR Image")262 img_thresh = gr.Slider(263 0.1, 0.9, value=0.5, step=0.05, label="Confidence Threshold"264 )265 img_btn = gr.Button("Detect Ships", variant="primary")266 with gr.Column():267 img_output = gr.Image(type="numpy", label="Detection Result")268 img_info = gr.Textbox(label="Detection Info", lines=10)269 270 img_btn.click(271 fn=detect_ships_image,272 inputs=[img_input, img_thresh],273 outputs=[img_output, img_info],274 )275 276 with gr.Tab("🎥 Video Detection"):277 gr.Markdown(278 "> ⚠️ CPU inference is slow. Short clips (< 30 s) are recommended."279 )280 with gr.Row():281 with gr.Column():282 vid_input = gr.Video(label="Upload SAR Video")283 vid_thresh = gr.Slider(284 0.1, 0.9, value=0.5, step=0.05, label="Confidence Threshold"285 )286 vid_btn = gr.Button("Detect Ships in Video", variant="primary")287 with gr.Column():288 vid_output = gr.Video(label="Detection Result Video")289 vid_info = gr.Textbox(label="Detection Summary", lines=8)290 291 vid_btn.click(292 fn=detect_ships_video,293 inputs=[vid_input, vid_thresh],294 outputs=[vid_output, vid_info],295 )296 297if __name__ == "__main__":298 demo.queue()299 demo.launch(server_name="0.0.0.0", server_port=7860) # NO share=True