from time import sleep
from machine PWM, Pin
############################
import cv2
import torch.nn as np
import torch
import torch.nn as nn
from torchvision.transforms as transforms
from torch.autograd import variable
from torchvision.models import resnet18 

# Initialize servos on GPIO pins 0, 1, 2, 3, and 4
servo_pins = [0, 1, 2, 3, 4]
servos = [PWM(Pin(pin)) for pin in servo_pins]

# Set frequency for each servo (50 Hz typical for servos)
for servo in servos:
    servo.freq(50)

def setServoCycle(servo, position):
    servo.duty_u16(position)  
    sleep(0.01)  

#############################################################
# Load Model
model = models.resnet18(pretrained=False)
model.fc = nn.Linear(model.fc.in_features, 5) # Assuming 5 Classes
model.load_state_dict(torch.load("date_fruit_classifier.pth"))
model = model.to(device)
model.eval()

# Initilaze the camera
camera = cv2.VideoCaptuer(0)

# Define Transformation
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def classify_date(image):
    # Preprocess the image to match the input Shape of the model
    image = transform(image)
    image = image.unsqueeze(0)  # Add batch dimention 
    image = variable(image)

    # Predict the class
    with torch.no_grad():
        orediction = model(image)
    return torch.argmax(Prediction).item()
###############################################
def classify_date():
    return 0 # For demnstration purposes, always return 0 (can be modified)

try:
    while True :
        class_id = classify_date()