TimCCG / app.py
syaikhipin's picture
Update app.py
abda3d8
# -*- coding: utf-8 -*-
"""Corrosion Excel.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/1XTrJK3G8Har_yZt83wasZX1LYD6DG1Sa
"""
#install dependencies
#!pip install roboflow
#!pip install gradio
#import library
import gradio as gr
from roboflow import Roboflow
import numpy as np
from PIL import Image
import requests
from io import BytesIO
import pandas as pd
import os
# Initialize Roboflow with your API key
rf = Roboflow(api_key="kKDoCn3ABT9AKeFQDCB4")
# Function to calculate the area of a polygon using the shoelace formula
def calculate_polygon_area(points):
n = len(points)
area = 0.0
for i in range(n):
x1, y1 = points[i]
x2, y2 = points[(i + 1) % n]
area += (x1 * y2 - x2 * y1)
return abs(area) / 2.0
# Function to process Roboflow prediction JSON and calculate corrosion areas
def calculate_corrosion_areas(json_data, unit="pixels", conversion_factor=1):
corrosion_areas = []
for prediction in json_data["predictions"]:
if prediction["class"] == "Corrosion":
points = [(point["x"], point["y"]) for point in prediction["points"]]
area = calculate_polygon_area(points)
if unit == "cm??":
area *= conversion_factor # Convert area from pixels to cm??
corrosion_areas.append(area)
total_corrosion_area = sum(corrosion_areas)
# Prepare output
result = {
"individual_areas": [f"{area} {unit}" for area in corrosion_areas],
"total_area": f"{total_corrosion_area} {unit}",
"recommendation": get_inspection_recommendation(total_corrosion_area)
}
return result
# Function to provide inspection recommendation based on total corrosion area
def get_inspection_recommendation(total_area):
if total_area < 1000:
return "No immediate inspection needed."
elif total_area < 5000:
return "Schedule an inspection in the next 6 months."
else:
return "Immediate inspection required."
# Define a Gradio interface to input a URL, run inference, and calculate corrosion areas
def url_infer_and_calculate(url, location, unit="pixels", conversion_factor=1):
try:
# Run inference using the Roboflow script
rf_project = rf.workspace().project("corrosion-instance-segmentation-sfcpc")
model = rf_project.version(3).model
prediction = model.predict(url)
# Ensure the response is properly formatted as JSON
prediction_json = prediction.json()
# Calculate corrosion areas from the Roboflow prediction
corrosion_areas = calculate_corrosion_areas(prediction_json, unit, float(conversion_factor))
# Download the image from the URL and convert it to a PIL Image
response = requests.get(url)
img = Image.open(BytesIO(response.content))
# Create a pandas DataFrame for reporting
df = pd.DataFrame([{'Number': index+1, 'URL': url, 'Location': location, 'corrosion_areas': corrosion_areas, 'Recommendation': corrosion_areas['recommendation']} for index in range(len(corrosion_areas))])
# Write DataFrame to local CSV file with index included immediately after creating it.
df.to_csv('Corrosion_Report.csv', index=False)
return img, corrosion_areas, prediction_json
except Exception as e:
return {"error": str(e)}
# Create a Gradio interface for URL input, inference, and corrosion area calculation
iface = gr.Interface(
fn=url_infer_and_calculate,
inputs=[
gr.inputs.Textbox(label="Enter the URL of an image"),
gr.inputs.Textbox(label="Enter the Location"),
gr.inputs.Dropdown(choices=["pixels", "cm??"], label="Area Unit"),
gr.inputs.Textbox(label="Conversion Factor")
],
outputs=[gr.outputs.Image(type="pil"), "json", "json"],
title="Tim CCG",
description="Enter the URL of an image to perform rust detection and calculate corrosion areas.",
)
# Launch the Gradio interface
iface.launch(debug=False, share=False)