Spaces:

hank1229
/

ML

Sleeping

ML / app.py

Update app.py

b20a292 verified over 1 year ago

3.85 kB

import tensorflow as tf
import joblib
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr

# 1. Load pretrained models
cnn_model = tf.keras.models.load_model('best_cnn_model.keras')   # Load CNN model
ann_model = tf.keras.models.load_model('best_ann_model.keras')   # Load ANN model

# Attempt to load linear regression models (assumed dict of four regressors)
try:
    lr_models = joblib.load('linear_regression_models.pkl')
except Exception as e:
    print(f"[Warning] Failed to load linear regression models: {e}")
    lr_models = None

# 2. Define prediction function with auto‑swap for manual bbox
def predict_bbox(model_choice, input_image, stop_x1, stop_y1, stop_x2, stop_y2):
    errors = []

    # Auto-swap manual bbox so that x0<=x1, y0<=y1
    mx0, mx1 = sorted([stop_x1, stop_x2])
    my0, my1 = sorted([stop_y1, stop_y2])

    # Preprocess image
    gray = input_image.convert('L')
    resized = gray.resize((128, 128))
    norm = np.array(resized) / 255.0

    # Model prediction branch
    if model_choice in ['CNN', 'ANN']:
        arr = norm.reshape(1, 128, 128, 1)
        if model_choice == 'CNN':
            pred = cnn_model.predict(arr)[0]
        else:
            pred = ann_model.predict(arr)[0]
    else:
        if lr_models is None:
            return "Error: Linear Regression models not loaded.", input_image
        flat = norm.flatten().reshape(1, -1)
        x1_p = lr_models['x1'].predict(flat)[0]
        y1_p = lr_models['y1'].predict(flat)[0]
        x2_p = lr_models['x2'].predict(flat)[0]
        y2_p = lr_models['y2'].predict(flat)[0]
        pred = np.array([x1_p, y1_p, x2_p, y2_p])

    # Scale predictions back to original image size
    orig_w, orig_h = input_image.size
    x1_p, y1_p, x2_p, y2_p = pred
    x0_p = int(min(x1_p, x2_p) * orig_w)
    x1_p = int(max(x1_p, x2_p) * orig_w)
    y0_p = int(min(y1_p, y2_p) * orig_h)
    y1_p = int(max(y1_p, y2_p) * orig_h)

    # Validate predicted bbox ordering
    if x1_p < x0_p:
        errors.append("Predicted bbox error: x2 < x1")
    if y1_p < y0_p:
        errors.append("Predicted bbox error: y2 < y1")

    # Create drawable copy
    img_draw = input_image.copy()
    draw = ImageDraw.Draw(img_draw)

    # Draw predicted bounding box in green (only if no critical error)
    if not errors:
        draw.rectangle([x0_p, y0_p, x1_p, y1_p], outline='green', width=2)
    else:
        # If prediction invalid, overlay error text
        err_text = "\n".join(errors)
        draw.text((10, 10), err_text, fill='red')

    # Draw manual Stop Sign box in red
    draw.rectangle([mx0, my0, mx1, my1], outline='red', width=2)

    # Prepare output text
    coords_text = (
        f"Model: {model_choice}\n"
        f"Predicted bbox: ({x0_p}, {y0_p}), ({x1_p}, {y1_p})\n"
        f"Manual Stop Sign bbox: ({mx0}, {my0}), ({mx1}, {my1})"
    )
    if errors:
        coords_text = "Errors in prediction:\n" + "\n".join(errors) + "\n\n" + coords_text

    return coords_text, img_draw

# 3. Build Gradio interface
iface = gr.Interface(
    fn=predict_bbox,
    inputs=[
        gr.Dropdown(choices=['CNN', 'ANN', 'Linear Regression'], label='Select Model'),
        gr.Image(type='pil', label='Upload Test Image'),
        gr.Number(label='Stop Sign x1 or x2'),
        gr.Number(label='Stop Sign y1 or y2'),
        gr.Number(label='Stop Sign x2 or x1'),
        gr.Number(label='Stop Sign y2 or y1'),
    ],
    outputs=[
        gr.Textbox(label='Results / Errors'),
        gr.Image(type='pil', label='Annotated Image'),
    ],
    title='Stop Sign Detection with Auto‑Swap Manual BBox',
    description=(
        'Manual bbox will auto‑swap if coords reversed. '
        'Predicted bbox in green; manual bbox in red. '
        'Prediction errors overlaid as text.'
    )
)

if __name__ == "__main__":
    iface.launch()

	import tensorflow as tf
	import joblib
	import numpy as np
	from PIL import Image, ImageDraw
	import gradio as gr

	# 1. Load pretrained models
	cnn_model = tf.keras.models.load_model('best_cnn_model.keras') # Load CNN model
	ann_model = tf.keras.models.load_model('best_ann_model.keras') # Load ANN model

	# Attempt to load linear regression models (assumed dict of four regressors)
	try:
	lr_models = joblib.load('linear_regression_models.pkl')
	except Exception as e:
	print(f"[Warning] Failed to load linear regression models: {e}")
	lr_models = None

	# 2. Define prediction function with auto‑swap for manual bbox
	def predict_bbox(model_choice, input_image, stop_x1, stop_y1, stop_x2, stop_y2):
	errors = []

	# Auto-swap manual bbox so that x0<=x1, y0<=y1
	mx0, mx1 = sorted([stop_x1, stop_x2])
	my0, my1 = sorted([stop_y1, stop_y2])

	# Preprocess image
	gray = input_image.convert('L')
	resized = gray.resize((128, 128))
	norm = np.array(resized) / 255.0

	# Model prediction branch
	if model_choice in ['CNN', 'ANN']:
	arr = norm.reshape(1, 128, 128, 1)
	if model_choice == 'CNN':
	pred = cnn_model.predict(arr)[0]
	else:
	pred = ann_model.predict(arr)[0]
	else:
	if lr_models is None:
	return "Error: Linear Regression models not loaded.", input_image
	flat = norm.flatten().reshape(1, -1)
	x1_p = lr_models['x1'].predict(flat)[0]
	y1_p = lr_models['y1'].predict(flat)[0]
	x2_p = lr_models['x2'].predict(flat)[0]
	y2_p = lr_models['y2'].predict(flat)[0]
	pred = np.array([x1_p, y1_p, x2_p, y2_p])

	# Scale predictions back to original image size
	orig_w, orig_h = input_image.size
	x1_p, y1_p, x2_p, y2_p = pred
	x0_p = int(min(x1_p, x2_p) * orig_w)
	x1_p = int(max(x1_p, x2_p) * orig_w)
	y0_p = int(min(y1_p, y2_p) * orig_h)
	y1_p = int(max(y1_p, y2_p) * orig_h)

	# Validate predicted bbox ordering
	if x1_p < x0_p:
	errors.append("Predicted bbox error: x2 < x1")
	if y1_p < y0_p:
	errors.append("Predicted bbox error: y2 < y1")

	# Create drawable copy
	img_draw = input_image.copy()
	draw = ImageDraw.Draw(img_draw)

	# Draw predicted bounding box in green (only if no critical error)
	if not errors:
	draw.rectangle([x0_p, y0_p, x1_p, y1_p], outline='green', width=2)
	else:
	# If prediction invalid, overlay error text
	err_text = "\n".join(errors)
	draw.text((10, 10), err_text, fill='red')

	# Draw manual Stop Sign box in red
	draw.rectangle([mx0, my0, mx1, my1], outline='red', width=2)

	# Prepare output text
	coords_text = (
	f"Model: {model_choice}\n"
	f"Predicted bbox: ({x0_p}, {y0_p}), ({x1_p}, {y1_p})\n"
	f"Manual Stop Sign bbox: ({mx0}, {my0}), ({mx1}, {my1})"
	)
	if errors:
	coords_text = "Errors in prediction:\n" + "\n".join(errors) + "\n\n" + coords_text

	return coords_text, img_draw

	# 3. Build Gradio interface
	iface = gr.Interface(
	fn=predict_bbox,
	inputs=[
	gr.Dropdown(choices=['CNN', 'ANN', 'Linear Regression'], label='Select Model'),
	gr.Image(type='pil', label='Upload Test Image'),
	gr.Number(label='Stop Sign x1 or x2'),
	gr.Number(label='Stop Sign y1 or y2'),
	gr.Number(label='Stop Sign x2 or x1'),
	gr.Number(label='Stop Sign y2 or y1'),
	],
	outputs=[
	gr.Textbox(label='Results / Errors'),
	gr.Image(type='pil', label='Annotated Image'),
	],
	title='Stop Sign Detection with Auto‑Swap Manual BBox',
	description=(
	'Manual bbox will auto‑swap if coords reversed. '
	'Predicted bbox in green; manual bbox in red. '
	'Prediction errors overlaid as text.'
	)
	)

	if __name__ == "__main__":
	iface.launch()