app.py

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('''
#       <p style="margin-bottom: 10px; font-size: 94%">
#         This is a demo for BRIA RMBG 1.4 that using
#         <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone. 
#       </p>
#     ''')
#     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('''
  <p style="margin-bottom: 10px; font-size: 94%">
    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
    <a href="https://huggingface.co/spaces/laitkor/product-image" target="_blank">Rootsraja internal product tool using an image matting model</a> as backbone. 
  </p>
''')
title = "Background Removal, Cropping, Resizing, Ruler and scaling"
description = r"""<a href="https://huggingface.co/spaces/laitkor/product-image" target="_blank">Rootsraja internal product tool<br> 
To use upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>.<br>
"""
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)