diff --git a/SKPhotoBrowser/SKPhoto.swift b/SKPhotoBrowser/SKPhoto.swift
index 5eace33..3baff81 100644
--- a/SKPhotoBrowser/SKPhoto.swift
+++ b/SKPhotoBrowser/SKPhoto.swift
@@ -11,6 +11,7 @@ import UIKit
 public protocol SKPhotoProtocol: NSObjectProtocol {
     var underlyingImage: UIImage! { get }
     var caption: String! { get }
+    var index: Int? { get set}
     func loadUnderlyingImageAndNotify()
     func checkCache()
 }
@@ -22,7 +23,8 @@ public class SKPhoto: NSObject, SKPhotoProtocol {
     public var photoURL: String!
     public var shouldCachePhotoURLImage: Bool = false
     public var caption: String!
-    
+    public var index: Int?
+
     override init() {
         super.init()
     }
diff --git a/SKPhotoBrowser/SKPhotoBrowser.swift b/SKPhotoBrowser/SKPhotoBrowser.swift
index c968454..6e937f6 100644
--- a/SKPhotoBrowser/SKPhotoBrowser.swift
+++ b/SKPhotoBrowser/SKPhotoBrowser.swift
@@ -14,7 +14,7 @@ import UIKit
     optional func willShowActionSheet(photoIndex: Int)
     optional func didDismissAtPageIndex(index: Int)
     optional func didDismissActionSheetWithButtonIndex(buttonIndex: Int, photoIndex: Int)
-    optional func didDeleted(notDeleted: [SKPhoto]) -> Void
+    optional func didDeleted(deletedIndex: [Int]) -> Void
 }
 
 public let SKPHOTO_LOADING_DID_END_NOTIFICATION = "photoLoadingDidEndNotification"
@@ -98,7 +98,7 @@ public class SKPhotoBrowser: UIViewController, UIScrollViewDelegate, UIActionShe
     var numberOfPhotos: Int {
         return photos.count
     }
-    
+    var deleted = [Int]()
     // MARK - Initializer
     required public init?(coder aDecoder: NSCoder) {
         super.init(coder: aDecoder)
@@ -251,6 +251,12 @@ public class SKPhotoBrowser: UIViewController, UIScrollViewDelegate, UIActionShe
     public override func viewWillAppear(animated: Bool) {
         super.viewWillAppear(true)
         reloadData()
+        
+        var i = 0
+        for photo : SKPhotoProtocol in photos {
+            photo.index = i
+            i = i + 1
+        }
     }
     
     public override func viewWillLayoutSubviews() {
@@ -627,7 +633,7 @@ public class SKPhotoBrowser: UIViewController, UIScrollViewDelegate, UIActionShe
         prepareForClosePhotoBrowser()
         dismissViewControllerAnimated(true) {
             self.delegate?.didDismissAtPageIndex?(self.currentPageIndex)
-            self.delegate?.didDeleted?(self.photos as! [SKPhoto])
+            self.delegate?.didDeleted?(self.deleted)
         }
     }
 
@@ -850,11 +856,15 @@ public class SKPhotoBrowser: UIViewController, UIScrollViewDelegate, UIActionShe
 
     // MARK: - Button
     public func deleteButtonPressed(sender: UIButton) {
-        photos.removeAtIndex(currentPageIndex)
-        if photos.count == 0 {
-            dismissPhotoBrowser()
-        } else {
+        let index = photos[currentPageIndex].index
+        deleted.append(index!)
+        print(deleted)
+        
+        if photos.count > 1 {
+            photos.removeAtIndex(currentPageIndex)
             reloadData()
+        } else {
+            dismissPhotoBrowser()
         }
     }