import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
import gradio as gr
from gradio_imageslider import ImageSlider
from briarmbg import BriaRMBG
import PIL
from PIL import Image, ImageDraw, ImageFont
from typing import Tuple
net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.to(device)
def draw_ruler(image, actual_height_in_inches = 13, dpi=300, unit="in", color="black"):
# Load the plant image in its original size
canvas = image
plant_width, plant_height = canvas.size
#actual_height_in_inches = 13 #only do the changes in this and generate the images ,,,,,,
#think about approximate height of plant suppose it is 8.3 inch so mention here 9 .and run the code
# Calculate pixels per inch based on the plant image's pixel height
pixels_per_inch = plant_height / actual_height_in_inches
# Create a blank canvas larger than the plant image to accommodate rulers
canvas_width = plant_width #+ 150 # 50 pixels extra for the left ruler
canvas_height = plant_height #+ 150 # 50 pixels extra for the bottom ruler
#canvas = Image.new('RGB', (canvas_width, canvas_height), 'white')
draw = ImageDraw.Draw(canvas)
font = ImageFont.load_default(size=60)
# for drawing scale value on vertical side ................ 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5...................so on
for i in range(0, int(plant_height / (pixels_per_inch * 0.5)) + 1):
# Calculate position and inch value
#y_position = canvas_height - 150 - int(i * pixels_per_inch * 0.5)
y_position = canvas_height - 0 - int(i * pixels_per_inch * 0.5)
inch_value = i * 0.5
# Draw tick mark and add text label every 1 inch for better readability
draw.line((0, y_position, 10, y_position), fill='black')
if inch_value.is_integer() and inch_value % 5 == 0: # Display labels only at whole numbers
draw.text((15, y_position - 5), f"{int(inch_value)} in", fill='black',font=font)
#else:
# draw.text((15, y_position - 5), f"{inch_value:.1f} in", fill='black',font=font)
# for drawing scale value on horizontal side ................ 0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5...................so on
for i in range(0, int(plant_width / (pixels_per_inch * 0.5)) + 1):
# Calculate position and inch value
#x_position = 150 + int(i * pixels_per_inch * 0.5)
x_position = 0 + int(i * pixels_per_inch * 0.5)
inch_value = i * 0.5
# Draw tick mark and add text label every 1 inch for better readability
draw.line((x_position, canvas_height - 40, x_position, canvas_height), fill='black')
if inch_value.is_integer() and inch_value % 5 == 0: # Display labels only at whole numbers
draw.text((x_position, canvas_height - 65), f"{int(inch_value)} in", fill='black',font=font)
#else:
#draw.text((x_position, canvas_height - 25), f"{inch_value:.1f} in", fill='black',font=font)
# Position to paste the plant image on the canvas (leaving space for rulers)
#plant_position = (150, canvas_height - plant_height - 150)
# Paste the original plant image onto the canvas without resizing
#canvas.paste(plant_image, plant_position)
# Save the final image with the plant and rulers
#canvas.save('image_with_plant_and_rulers_high_quality_ONE.png')
#print("The image with the plant and rulers has been saved at original quality.")
return canvas;
def draw_ruler0(image, dpi=300, unit="in", color="black"):
# Define the pot dimensions and distance from camera
pot_diameter_inches = 4
pot_height_inches = 6
distance_from_camera_feet = 0.5
zoom_level = 1 # Zoom level
# Convert distance from feet to inches
distance_from_camera_inches = distance_from_camera_feet * 12
# Calculate the scale factor (assuming a simple pinhole camera model)
# Scale factor = (actual height of the object) / (height of the object in the image)
# Adjust for zoom level
scale_factor = (pot_height_inches / distance_from_camera_inches) * zoom_level
print ("scale factor:",scale_factor)
# Draw the scale on the image
draw = ImageDraw.Draw(image)
font = ImageFont.load_default()
# Define the position and size of the scale
scale_length_pixels = int(pot_height_inches / scale_factor)
scale_position = (50, 50) # Adjust as needed
# Draw the scale line
draw.line((scale_position[0], scale_position[1], scale_position[0], scale_position[1] + scale_length_pixels), fill="red", width=5)
print("line:",(scale_position[0], scale_position[1], scale_position[0], scale_position[1] + scale_length_pixels))
# Add text label for the scale
draw.text((scale_position[0] + 10, scale_position[1] + scale_length_pixels // 2), f"{pot_height_inches} inches", fill="red", font=font)
return image
def draw_ruler1(image, dpi=300, unit="in", color="black"):
"""
Draws a ruler on the top and left edges of an image.
Args:
image: The PIL Image object to draw on.
dpi: The DPI of the image (default 300).
unit: The unit of measurement (default "in").
color: The color of the ruler (default "black").
"""
#dpi = image.info.get("dpi", (72, 72))
#dpi = dpi[0] if isinstance(dpi, tuple) else dpi
# Define the pot dimensions and distance from camera
pot_diameter_inches = 4
pot_height_inches = 6
distance_from_camera_feet = 3
zoom_level = 0.5 # Zoom level
# Convert distance from feet to inches
distance_from_camera_inches = distance_from_camera_feet * 12
# Calculate the scale factor (assuming a simple pinhole camera model)
# Scale factor = (actual height of the object) / (height of the object in the image)
# Adjust for zoom level
scale_factor = (pot_height_inches / distance_from_camera_inches) * zoom_level
scale_length_pixels = int(pot_height_inches / scale_factor)
print ("scale factor:",scale_factor,"scale_length_pixels:",scale_length_pixels)
draw = ImageDraw.Draw(image)
width, height = image.size
font = ImageFont.load_default(size=60)
# Draw top ruler
#for i in range(0, width, int(dpi)):
for i in range(0, width, scale_length_pixels):
x = i
draw.line((x, 0, x, 10), fill=color)
if i % int(scale_length_pixels) == 0:
if scale_length_pixels< 70 :
draw.text((x, 12), str(i // scale_length_pixels), fill=color, font=font)
if scale_length_pixels > 70 and i % 5 == 0:
draw.text((x, 12), str(i // scale_length_pixels), fill=color, font=font)
# Draw left ruler
for i in range(0, height, int(dpi)):
y = i
draw.line((0, y, 10, y), fill=color)
if i % int(dpi) == 0:
if i != 0:
draw.text((12, y), str(i // dpi), fill=color, font=font)
draw.text((100, y-60), "inch", fill=color, font=font)
return image
def mm_to_pixels(mm, dpi):
if dpi is None:
dpi = 72 # Default DPI value
if mm is None:
mm = 0 # Default mm value
return int(mm * dpi / 25.4)
def crop_image(image, left_mm, top_mm, right_mm, bottom_mm):
# Open the image file
#img = Image.open(image)
img = image
print("img.width:",img.width,"img.height:",img.height)
# Get the image DPI (dots per inch)
#dpi = img.info.get('dpi', (72, 72))[0] # Default to 72 DPI if not available
# Get the DPI value as a single number
dpi = img.info.get("dpi", (72, 72))
dpi = dpi[0] if isinstance(dpi, tuple) else dpi
print ("dpi:",dpi)
# Convert mm to pixels
left = mm_to_pixels(left_mm, dpi)
top = mm_to_pixels(top_mm, dpi)
right = mm_to_pixels(right_mm, dpi)
bottom = mm_to_pixels(bottom_mm, dpi)
print("left:",left,"top:",top,"right:",right,"bottom:",bottom)
if img.width - right <= left :
right=img.width-left-1
# Crop the image
cropped_img = img.crop((left, top, img.width - right, img.height - bottom))
return cropped_img
def resize_image(image):
image = image.convert('RGB')
model_input_size = (1024, 1024)
image = image.resize(model_input_size, Image.BILINEAR)
return image
def process(image,left_mm, top_mm, right_mm, bottom_mm, ruler, resize,actual_height_in_inches,background):
# prepare input
orig_image = Image.fromarray(image)
print("orig size:",orig_image.size)
w,h = orig_im_size = orig_image.size
image = resize_image(orig_image)
new_im=image;
if background:
# remove background
im_np = np.array(image)
im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
im_tensor = torch.unsqueeze(im_tensor,0)
im_tensor = torch.divide(im_tensor,255.0)
im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
if torch.cuda.is_available():
im_tensor=im_tensor.cuda()
#inference
result=net(im_tensor)
# post process
result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
ma = torch.max(result)
mi = torch.min(result)
result = (result-mi)/(ma-mi)
# image to pil
im_array = (result*255).cpu().data.numpy().astype(np.uint8)
pil_im = Image.fromarray(np.squeeze(im_array))
# paste the mask on the original image
new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
new_im.paste(orig_image, mask=pil_im)
# new_orig_image = orig_image.convert('RGBA')
else:
new_im=image;
new_im= crop_image(new_im, left_mm, top_mm, right_mm, bottom_mm)
if ruler:
new_im = draw_ruler(new_im,actual_height_in_inches)
if resize is not None and resize > 0:
new_im = new_im.resize( [int(resize * s) for s in new_im.size],Image.Resampling.LANCZOS )
print("resize:",new_im.size)
return new_im
# return [new_orig_image, new_im]
# block = gr.Blocks().queue()
# with block:
# gr.Markdown("## BRIA RMBG 1.4")
# gr.HTML('''
#
# This is a demo for BRIA RMBG 1.4 that using
# BRIA RMBG-1.4 image matting model as backbone.
#
# ''')
# with gr.Row():
# with gr.Column():
# input_image = gr.Image(sources=None, type="pil") # None for upload, ctrl+v and webcam
# # input_image = gr.Image(sources=None, type="numpy") # None for upload, ctrl+v and webcam
# run_button = gr.Button(value="Run")
# with gr.Column():
# result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", columns=[1], height='auto')
# ips = [input_image]
# run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
# block.launch(debug = True)
# block = gr.Blocks().queue()
gr.Markdown("## BRIA RMBG 1.4")
gr.HTML('''
This app is leveraging the RMBG v1.4 model which is developed on the IS-Net enhanced with unique training scheme and proprietary dataset.
These modifications significantly improve the model’s accuracy and effectiveness in diverse image-processing scenarios.
This is an internal tool being tested for rootsraja.in product images at
Rootsraja internal product tool using an image matting model as backbone.
''')
title = "Background Removal, Cropping, Resizing, Ruler and scaling"
description = r"""Rootsraja internal product tool
To use upload your image and wait. Read more at model card briaai/RMBG-1.4.
"""
examples = [['./input.jpg'],]
# output = ImageSlider(position=0.5,label='Image without background', type="pil", show_download_button=True)
# demo = gr.Interface(fn=process,inputs="image", outputs=output, examples=examples, title=title, description=description)
demo = gr.Interface(fn=process,
inputs=[
gr.Image(type="numpy", label="Upload Image"),
gr.Number(label="Left Crop (mm)",value=0),
gr.Number(label="Top Crop (mm)",value=0),
gr.Number(label="Right Crop (mm)",value=0),
gr.Number(label="Bottom Crop (mm)",value=0),
gr.Checkbox(label="Ruler!"),
gr.Number(label="Resize (between 0.1 and 0.9)",value=0),
gr.Number(label="Plant height in inches",value=0),
gr.Checkbox(label="Remove Background?")
],
outputs="image", examples=examples, title=title, description=description)
if __name__ == "__main__":
demo.launch(share=False)