LPX
commited on
Commit
·
932e7b4
1
Parent(s):
6efb635
major(huge refactoring)
Browse files- Added entries to .gitignore for Python cache files and directories.
- Updated README.md title, description, and SDK version.
- Changed app file reference from app.py to app_mcp.py.
- .gitignore +3 -1
- README.md +5 -5
- app_mcp.py +326 -0
- forensics/registry.py +13 -0
.gitignore
CHANGED
@@ -2,4 +2,6 @@
|
|
2 |
*goat.py
|
3 |
.vscode
|
4 |
*onnx.py
|
5 |
-
./models/*
|
|
|
|
|
|
2 |
*goat.py
|
3 |
.vscode
|
4 |
*onnx.py
|
5 |
+
./models/*
|
6 |
+
forensics/__pycache__/*
|
7 |
+
*.cpython-311.pyc
|
README.md
CHANGED
@@ -1,11 +1,12 @@
|
|
1 |
---
|
2 |
-
title:
|
3 |
-
|
|
|
4 |
colorFrom: yellow
|
5 |
colorTo: yellow
|
6 |
sdk: gradio
|
7 |
-
sdk_version: 5.
|
8 |
-
app_file:
|
9 |
pinned: true
|
10 |
models:
|
11 |
- aiwithoutborders-xyz/OpenSight-CommunityForensics-Deepfake-ViT
|
@@ -13,7 +14,6 @@ models:
|
|
13 |
- haywoodsloan/ai-image-detector-deploy
|
14 |
- cmckinle/sdxl-flux-detector
|
15 |
- Organika/sdxl-detector
|
16 |
-
- prithivMLmods/Deepfake-Detection-Exp-02-22
|
17 |
license: mit
|
18 |
---
|
19 |
|
|
|
1 |
---
|
2 |
+
title: Deepfake Detection & Forensics Tools
|
3 |
+
description: MCP Server for Deepfake Detection & Digital Forensics Tools
|
4 |
+
emoji: 🚑
|
5 |
colorFrom: yellow
|
6 |
colorTo: yellow
|
7 |
sdk: gradio
|
8 |
+
sdk_version: 5.33.0
|
9 |
+
app_file: app_mcp.py
|
10 |
pinned: true
|
11 |
models:
|
12 |
- aiwithoutborders-xyz/OpenSight-CommunityForensics-Deepfake-ViT
|
|
|
14 |
- haywoodsloan/ai-image-detector-deploy
|
15 |
- cmckinle/sdxl-flux-detector
|
16 |
- Organika/sdxl-detector
|
|
|
17 |
license: mit
|
18 |
---
|
19 |
|
app_mcp.py
ADDED
@@ -0,0 +1,326 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from typing import Literal
|
3 |
+
import spaces
|
4 |
+
import gradio as gr
|
5 |
+
import modelscope_studio.components.antd as antd
|
6 |
+
import modelscope_studio.components.antdx as antdx
|
7 |
+
import modelscope_studio.components.base as ms
|
8 |
+
from transformers import pipeline, AutoImageProcessor, SwinForImageClassification, Swinv2ForImageClassification, AutoFeatureExtractor, AutoModelForImageClassification
|
9 |
+
from torchvision import transforms
|
10 |
+
import torch
|
11 |
+
from PIL import Image
|
12 |
+
import numpy as np
|
13 |
+
import io
|
14 |
+
import logging
|
15 |
+
from utils.utils import softmax, augment_image, convert_pil_to_bytes
|
16 |
+
from utils.gradient import gradient_processing
|
17 |
+
from utils.minmax import preprocess as minmax_preprocess
|
18 |
+
from utils.ela import genELA as ELA
|
19 |
+
from forensics.registry import register_model, MODEL_REGISTRY
|
20 |
+
|
21 |
+
|
22 |
+
# Configure logging
|
23 |
+
logging.basicConfig(level=logging.INFO)
|
24 |
+
logger = logging.getLogger(__name__)
|
25 |
+
|
26 |
+
|
27 |
+
# Ensure using GPU if available
|
28 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
29 |
+
|
30 |
+
header_style = {
|
31 |
+
"textAlign": 'center',
|
32 |
+
"color": '#fff',
|
33 |
+
"height": 64,
|
34 |
+
"paddingInline": 48,
|
35 |
+
"lineHeight": '64px',
|
36 |
+
"backgroundColor": '#4096ff',
|
37 |
+
}
|
38 |
+
|
39 |
+
content_style = {
|
40 |
+
"textAlign": 'center',
|
41 |
+
"minHeight": 120,
|
42 |
+
"lineHeight": '120px',
|
43 |
+
"color": '#fff',
|
44 |
+
"backgroundColor": '#0958d9',
|
45 |
+
}
|
46 |
+
|
47 |
+
sider_style = {
|
48 |
+
"textAlign": 'center',
|
49 |
+
"lineHeight": '120px',
|
50 |
+
"color": '#fff',
|
51 |
+
"backgroundColor": '#1677ff',
|
52 |
+
}
|
53 |
+
|
54 |
+
footer_style = {
|
55 |
+
"textAlign": 'center',
|
56 |
+
"color": '#fff',
|
57 |
+
"backgroundColor": '#4096ff',
|
58 |
+
}
|
59 |
+
|
60 |
+
layout_style = {
|
61 |
+
"borderRadius": 8,
|
62 |
+
"overflow": 'hidden',
|
63 |
+
"width": 'calc(100% - 8px)',
|
64 |
+
"maxWidth": 'calc(100% - 8px)',
|
65 |
+
}
|
66 |
+
# Model paths and class names
|
67 |
+
MODEL_PATHS = {
|
68 |
+
"model_1": "haywoodsloan/ai-image-detector-deploy",
|
69 |
+
"model_2": "Heem2/AI-vs-Real-Image-Detection",
|
70 |
+
"model_3": "Organika/sdxl-detector",
|
71 |
+
"model_4": "cmckinle/sdxl-flux-detector_v1.1",
|
72 |
+
"model_5": "prithivMLmods/Deep-Fake-Detector-v2-Model",
|
73 |
+
"model_5b": "prithivMLmods/Deepfake-Detection-Exp-02-22",
|
74 |
+
"model_6": "ideepankarsharma2003/AI_ImageClassification_MidjourneyV6_SDXL",
|
75 |
+
"model_7": "date3k2/vit-real-fake-classification-v4"
|
76 |
+
}
|
77 |
+
|
78 |
+
CLASS_NAMES = {
|
79 |
+
"model_1": ['artificial', 'real'],
|
80 |
+
"model_2": ['AI Image', 'Real Image'],
|
81 |
+
"model_3": ['AI', 'Real'],
|
82 |
+
"model_4": ['AI', 'Real'],
|
83 |
+
"model_5": ['Realism', 'Deepfake'],
|
84 |
+
"model_5b": ['Real', 'Deepfake'],
|
85 |
+
"model_6": ['ai_gen', 'human'],
|
86 |
+
"model_7": ['Fake', 'Real'],
|
87 |
+
|
88 |
+
}
|
89 |
+
|
90 |
+
def preprocess_resize_256(image):
|
91 |
+
if image.mode != 'RGB':
|
92 |
+
image = image.convert('RGB')
|
93 |
+
return transforms.Resize((256, 256))(image)
|
94 |
+
|
95 |
+
def preprocess_resize_224(image):
|
96 |
+
if image.mode != 'RGB':
|
97 |
+
image = image.convert('RGB')
|
98 |
+
return transforms.Resize((224, 224))(image)
|
99 |
+
|
100 |
+
def postprocess_pipeline(prediction, class_names):
|
101 |
+
# Assumes HuggingFace pipeline output
|
102 |
+
return {pred['label']: pred['score'] for pred in prediction}
|
103 |
+
|
104 |
+
def postprocess_logits(outputs, class_names):
|
105 |
+
# Assumes model output with logits
|
106 |
+
logits = outputs.logits.cpu().numpy()[0]
|
107 |
+
probabilities = softmax(logits)
|
108 |
+
return {class_names[i]: probabilities[i] for i in range(len(class_names))}
|
109 |
+
|
110 |
+
# Load and register models (example for two models)
|
111 |
+
image_processor_1 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_1"], use_fast=True)
|
112 |
+
model_1 = Swinv2ForImageClassification.from_pretrained(MODEL_PATHS["model_1"]).to(device)
|
113 |
+
clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
|
114 |
+
register_model(
|
115 |
+
"model_1",
|
116 |
+
clf_1,
|
117 |
+
preprocess_resize_256,
|
118 |
+
postprocess_pipeline,
|
119 |
+
CLASS_NAMES["model_1"]
|
120 |
+
)
|
121 |
+
|
122 |
+
clf_2 = pipeline("image-classification", model=MODEL_PATHS["model_2"], device=device)
|
123 |
+
register_model(
|
124 |
+
"model_2",
|
125 |
+
clf_2,
|
126 |
+
preprocess_resize_224,
|
127 |
+
postprocess_pipeline,
|
128 |
+
CLASS_NAMES["model_2"]
|
129 |
+
)
|
130 |
+
|
131 |
+
# Register remaining models
|
132 |
+
feature_extractor_3 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_3"], device=device)
|
133 |
+
model_3 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_3"]).to(device)
|
134 |
+
def preprocess_256(image):
|
135 |
+
if image.mode != 'RGB':
|
136 |
+
image = image.convert('RGB')
|
137 |
+
return transforms.Resize((256, 256))(image)
|
138 |
+
def postprocess_logits_model3(outputs, class_names):
|
139 |
+
logits = outputs.logits.cpu().numpy()[0]
|
140 |
+
probabilities = softmax(logits)
|
141 |
+
return {class_names[i]: probabilities[i] for i in range(len(class_names))}
|
142 |
+
def model3_infer(image):
|
143 |
+
inputs = feature_extractor_3(image, return_tensors="pt").to(device)
|
144 |
+
with torch.no_grad():
|
145 |
+
outputs = model_3(**inputs)
|
146 |
+
return outputs
|
147 |
+
register_model(
|
148 |
+
"model_3",
|
149 |
+
model3_infer,
|
150 |
+
preprocess_256,
|
151 |
+
postprocess_logits_model3,
|
152 |
+
CLASS_NAMES["model_3"]
|
153 |
+
)
|
154 |
+
|
155 |
+
feature_extractor_4 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_4"], device=device)
|
156 |
+
model_4 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_4"]).to(device)
|
157 |
+
def model4_infer(image):
|
158 |
+
inputs = feature_extractor_4(image, return_tensors="pt").to(device)
|
159 |
+
with torch.no_grad():
|
160 |
+
outputs = model_4(**inputs)
|
161 |
+
return outputs
|
162 |
+
def postprocess_logits_model4(outputs, class_names):
|
163 |
+
logits = outputs.logits.cpu().numpy()[0]
|
164 |
+
probabilities = softmax(logits)
|
165 |
+
return {class_names[i]: probabilities[i] for i in range(len(class_names))}
|
166 |
+
register_model(
|
167 |
+
"model_4",
|
168 |
+
model4_infer,
|
169 |
+
preprocess_256,
|
170 |
+
postprocess_logits_model4,
|
171 |
+
CLASS_NAMES["model_4"]
|
172 |
+
)
|
173 |
+
|
174 |
+
clf_5 = pipeline("image-classification", model=MODEL_PATHS["model_5"], device=device)
|
175 |
+
register_model(
|
176 |
+
"model_5",
|
177 |
+
clf_5,
|
178 |
+
preprocess_resize_224,
|
179 |
+
postprocess_pipeline,
|
180 |
+
CLASS_NAMES["model_5"]
|
181 |
+
)
|
182 |
+
|
183 |
+
clf_5b = pipeline("image-classification", model=MODEL_PATHS["model_5b"], device=device)
|
184 |
+
register_model(
|
185 |
+
"model_5b",
|
186 |
+
clf_5b,
|
187 |
+
preprocess_resize_224,
|
188 |
+
postprocess_pipeline,
|
189 |
+
CLASS_NAMES["model_5b"]
|
190 |
+
)
|
191 |
+
|
192 |
+
image_processor_6 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_6"], use_fast=True)
|
193 |
+
model_6 = SwinForImageClassification.from_pretrained(MODEL_PATHS["model_6"]).to(device)
|
194 |
+
clf_6 = pipeline(model=model_6, task="image-classification", image_processor=image_processor_6, device=device)
|
195 |
+
register_model(
|
196 |
+
"model_6",
|
197 |
+
clf_6,
|
198 |
+
preprocess_resize_224,
|
199 |
+
postprocess_pipeline,
|
200 |
+
CLASS_NAMES["model_6"]
|
201 |
+
)
|
202 |
+
|
203 |
+
image_processor_7 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_7"], use_fast=True)
|
204 |
+
model_7 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_7"]).to(device)
|
205 |
+
clf_7 = pipeline(model=model_7, task="image-classification", image_processor=image_processor_7, device=device)
|
206 |
+
register_model(
|
207 |
+
"model_7",
|
208 |
+
clf_7,
|
209 |
+
preprocess_resize_224,
|
210 |
+
postprocess_pipeline,
|
211 |
+
CLASS_NAMES["model_7"]
|
212 |
+
)
|
213 |
+
|
214 |
+
# Generic inference function
|
215 |
+
|
216 |
+
def infer(image: Image.Image, model_id: str, confidence_threshold: float = 0.75) -> dict:
|
217 |
+
entry = MODEL_REGISTRY[model_id]
|
218 |
+
img = entry.preprocess(image)
|
219 |
+
try:
|
220 |
+
result = entry.model(img)
|
221 |
+
result = entry.postprocess(result, entry.class_names)
|
222 |
+
# Add confidence threshold logic if needed
|
223 |
+
return result
|
224 |
+
except Exception as e:
|
225 |
+
return {"error": str(e)}
|
226 |
+
|
227 |
+
# Update predict_image to use all registered models in order
|
228 |
+
|
229 |
+
def predict_image(img, confidence_threshold):
|
230 |
+
model_ids = [
|
231 |
+
"model_1", "model_2", "model_3", "model_4", "model_5", "model_5b", "model_6", "model_7"
|
232 |
+
]
|
233 |
+
results = [infer(img, model_id, confidence_threshold) for model_id in model_ids]
|
234 |
+
return img, results
|
235 |
+
|
236 |
+
# Update predict_image_with_json to return results as a list of dicts
|
237 |
+
|
238 |
+
def predict_image_with_json(img, confidence_threshold, augment_methods, rotate_degrees, noise_level, sharpen_strength):
|
239 |
+
if augment_methods:
|
240 |
+
img_pil, _ = augment_image(img, augment_methods, rotate_degrees, noise_level, sharpen_strength)
|
241 |
+
else:
|
242 |
+
img_pil = img
|
243 |
+
img_pil, results = predict_image(img_pil, confidence_threshold)
|
244 |
+
img_np = np.array(img_pil) # Convert PIL Image to NumPy array
|
245 |
+
img_np_og = np.array(img) # Convert PIL Image to NumPy array
|
246 |
+
|
247 |
+
gradient_image = gradient_processing(img_np) # Added gradient processing
|
248 |
+
minmax_image = minmax_preprocess(img_np) # Added MinMax processing
|
249 |
+
|
250 |
+
# First pass - standard analysis
|
251 |
+
ela1 = ELA(img_np_og, quality=75, scale=50, contrast=20, linear=False, grayscale=True)
|
252 |
+
|
253 |
+
# Second pass - enhanced visibility
|
254 |
+
ela2 = ELA(img_np_og, quality=75, scale=75, contrast=25, linear=False, grayscale=True)
|
255 |
+
ela3 = ELA(img_np_og, quality=75, scale=75, contrast=25, linear=False, grayscale=False)
|
256 |
+
|
257 |
+
forensics_images = [img_pil, ela1, ela2, ela3, gradient_image, minmax_image]
|
258 |
+
|
259 |
+
return img_pil, forensics_images, results
|
260 |
+
|
261 |
+
with gr.Blocks(css="#post-gallery { overflow: hidden !important;} .grid-wrap{ overflow-y: hidden !important;} .ms-gr-ant-welcome-icon{ height:unset !important;} .tabs{margin-top:10px;}") as iface:
|
262 |
+
with ms.Application() as app:
|
263 |
+
with antd.ConfigProvider():
|
264 |
+
antdx.Welcome(
|
265 |
+
icon=
|
266 |
+
"https://cdn-avatars.huggingface.co/v1/production/uploads/639daf827270667011153fbc/WpeSFhuB81DY-1TjNUmV_.png",
|
267 |
+
title="Welcome to Project OpenSight",
|
268 |
+
description=
|
269 |
+
"The OpenSight aims to be an open-source SOTA generated image detection model. This HF Space is not only an introduction but a educational playground for the public to evaluate and challenge current open source models. **Space will be upgraded shortly; inference on all 6 models should take about 1.2~ seconds.** "
|
270 |
+
)
|
271 |
+
with gr.Tab("👀 Detection Models Eval / Playground"):
|
272 |
+
gr.Markdown("# Open Source Detection Models Found on the Hub\n\n - **Space will be upgraded shortly;** inference on all 6 models should take about 1.2~ seconds once we're back on CUDA.\n - The **Community Forensics** mother of all detection models is now available for inference, head to the middle tab above this.\n - Lots of exciting things coming up, stay tuned!")
|
273 |
+
|
274 |
+
with gr.Row():
|
275 |
+
with gr.Column(scale=1):
|
276 |
+
image_input = gr.Image(label="Upload Image to Analyze", sources=['upload', 'webcam'], type='pil')
|
277 |
+
with gr.Accordion("Settings (Optional)", open=False, elem_id="settings_accordion"):
|
278 |
+
augment_checkboxgroup = gr.CheckboxGroup(["rotate", "add_noise", "sharpen"], label="Augmentation Methods")
|
279 |
+
rotate_slider = gr.Slider(0, 45, value=2, step=1, label="Rotate Degrees", visible=False)
|
280 |
+
noise_slider = gr.Slider(0, 50, value=4, step=1, label="Noise Level", visible=False)
|
281 |
+
sharpen_slider = gr.Slider(0, 50, value=11, step=1, label="Sharpen Strength", visible=False)
|
282 |
+
confidence_slider = gr.Slider(0.0, 1.0, value=0.75, step=0.05, label="Confidence Threshold")
|
283 |
+
inputs = [image_input, confidence_slider, augment_checkboxgroup, rotate_slider, noise_slider, sharpen_slider]
|
284 |
+
predict_button = gr.Button("Predict")
|
285 |
+
augment_button = gr.Button("Augment & Predict")
|
286 |
+
image_output = gr.Image(label="Processed Image", visible=False)
|
287 |
+
|
288 |
+
|
289 |
+
with gr.Column(scale=2):
|
290 |
+
# Use Gradio-native Dataframe to display results
|
291 |
+
results_table = gr.Dataframe(label="Model Predictions", headers=None, datatype="auto")
|
292 |
+
forensics_gallery = gr.Gallery(label="Post Processed Images", visible=True, columns=[4], rows=[2], container=False, height="auto", object_fit="contain", elem_id="post-gallery")
|
293 |
+
|
294 |
+
outputs = [image_output, forensics_gallery, results_table]
|
295 |
+
|
296 |
+
# Show/hide rotate slider based on selected augmentation method
|
297 |
+
augment_checkboxgroup.change(lambda methods: gr.update(visible="rotate" in methods), inputs=[augment_checkboxgroup], outputs=[rotate_slider])
|
298 |
+
augment_checkboxgroup.change(lambda methods: gr.update(visible="add_noise" in methods), inputs=[augment_checkboxgroup], outputs=[noise_slider])
|
299 |
+
augment_checkboxgroup.change(lambda methods: gr.update(visible="sharpen" in methods), inputs=[augment_checkboxgroup], outputs=[sharpen_slider])
|
300 |
+
|
301 |
+
predict_button.click(
|
302 |
+
fn=predict_image_with_json,
|
303 |
+
inputs=inputs,
|
304 |
+
outputs=outputs
|
305 |
+
)
|
306 |
+
augment_button.click( # Connect Augment button to the function
|
307 |
+
fn=predict_image_with_json,
|
308 |
+
inputs=[
|
309 |
+
image_input,
|
310 |
+
confidence_slider,
|
311 |
+
gr.CheckboxGroup(["rotate", "add_noise", "sharpen"], value=["rotate", "add_noise", "sharpen"], visible=False), # Default values
|
312 |
+
rotate_slider,
|
313 |
+
noise_slider,
|
314 |
+
sharpen_slider
|
315 |
+
],
|
316 |
+
outputs=outputs
|
317 |
+
)
|
318 |
+
with gr.Tab("👑 Community Forensics Preview"):
|
319 |
+
temp_space = gr.load("aiwithoutborders-xyz/OpenSight-Community-Forensics-Preview", src="spaces")
|
320 |
+
# preview # no idea if this will work
|
321 |
+
with gr.Tab("🥇 Leaderboard"):
|
322 |
+
gr.Markdown("# AI Generated / Deepfake Detection Models Leaderboard: Soon™")
|
323 |
+
|
324 |
+
|
325 |
+
# Launch the interface
|
326 |
+
iface.launch()
|
forensics/registry.py
ADDED
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Callable, Dict, Any, List
|
2 |
+
|
3 |
+
class ModelEntry:
|
4 |
+
def __init__(self, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str]):
|
5 |
+
self.model = model
|
6 |
+
self.preprocess = preprocess
|
7 |
+
self.postprocess = postprocess
|
8 |
+
self.class_names = class_names
|
9 |
+
|
10 |
+
MODEL_REGISTRY: Dict[str, ModelEntry] = {}
|
11 |
+
|
12 |
+
def register_model(model_id: str, model: Any, preprocess: Callable, postprocess: Callable, class_names: List[str]):
|
13 |
+
MODEL_REGISTRY[model_id] = ModelEntry(model, preprocess, postprocess, class_names)
|