import SwiftUI
import Vision

struct ContentView: View {
   @State private var observations: SaliencyImageObservation?

   var body: some View {
      VStack {
         Button("Detect") {
            Task {
               await detectObjects()
            }
         }.buttonStyle(.borderedProminent)

         Canvas { context, size in
            // Calculate size and draw image
            let image = UIImage(named: "picturedog")!
            let width = size.width
            let height = image.size.height * width / image.size.width
            let imageFrame = CGRect(origin: .zero, size: CGSize(width: width, height: height))
            context.draw(Image(uiImage: image), in: imageFrame)

            // Calculate size and draw bounding box
            if let salientObjects = observations?.salientObjects {
               for object in salientObjects {
                  let rect = object.boundingBox
                  let boxFrame = rect.toImageCoordinates(imageFrame.size, origin: .upperLeft)
                  context.stroke(Rectangle().path(in: boxFrame), with: .color(.green), lineWidth: 5)
               }
            }
         }
         Spacer()
      }
      .padding()
   }
   func detectObjects() async {
      let imageURL = Bundle.main.url(forResource: "picturedog", withExtension: "png")
      if let imageURL = imageURL {
         do {
            let request = GenerateAttentionBasedSaliencyImageRequest()
            observations = try await request.perform(on: imageURL)
         } catch {
            print("Error performing the request: \(error)")
         }
      }
   }
}